#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
QNX系统进程监控工具

该脚本用于持续监控QNX系统中特定进程的CPU和内存使用情况，
通过ADB和telnet自动连接到设备并获取相关信息。
"""

import subprocess
import time
import re
import csv
import os
import shutil
import platform
from datetime import datetime

# 根据操作系统选择适当的库
if platform.system() == 'Windows':
    try:
        import wexpect

        # 创建一个别名，使代码兼容
        pexpect = wexpect
        IS_WINDOWS = True
        # Windows系统下，增加调试输出
        DEBUG_MODE = True


        # wexpect库没有exceptions模块，需要自定义异常处理
        class WexpectExceptionWrapper:
            class ExceptionPexpect(Exception):
                pass


        pexpect.exceptions = WexpectExceptionWrapper
    except ImportError:
        print("错误: 在Windows系统上需要安装wexpect库。")
        print("请运行: pip install wexpect")
        exit(1)
else:
    import pexpect

    IS_WINDOWS = False
    DEBUG_MODE = False

#################################################
################ 配置参数区域 ####################
#################################################

# 要监控的进程名称
PROCESS_NAME = "logservice"

# 数据采集间隔(秒)
INTERVAL = 5

# 设备类型选择
# "device1": 需要密码验证的设备 (默认)
# "device2": 只需账号不需密码的设备 (执行adb root后连接)
DEVICE_TYPE = "device1"

# 设备1配置 (需要密码验证)
DEVICE1_QNX_IP = "192.168.8.1"
DEVICE1_USERNAME = "root"
DEVICE1_PASSWORD = "sv2655888"

# 设备2配置 (只需账号无需密码)
DEVICE2_QNX_IP = "172.31.207.204"
DEVICE2_USERNAME = "root"

# 输出文件路径，留空则自动生成文件名
OUTPUT_FILE = ""  # 例如: "logservice_monitor.csv"

# ADB路径设置
# 如果为空，将自动查找adb路径；否则使用指定的路径
# 例如: "/Users/username/Library/Android/sdk/platform-tools/adb"
ADB_PATH = ""


#################################################
################ 程序代码区域 ####################
#################################################

def find_adb_path():
    """
    查找adb命令的路径

    Returns:
        adb命令的完整路径，如果未找到则返回None
    """
    # 首先检查是否已经设置了ADB_PATH
    if ADB_PATH and os.path.exists(ADB_PATH) and os.access(ADB_PATH, os.X_OK):
        return ADB_PATH

    # 检查环境变量中的adb
    adb_from_path = shutil.which('adb')
    if adb_from_path:
        return adb_from_path

    # 检查常见的adb安装路径
    common_paths = [
        # macOS 常见路径
        '/usr/local/bin/adb',
        os.path.expanduser('~/Library/Android/sdk/platform-tools/adb'),
        '/Applications/Android Studio.app/Contents/sdk/platform-tools/adb',
        # Linux 常见路径
        '/usr/bin/adb',
        '/opt/android-sdk/platform-tools/adb',
        # Windows 常见路径
        'C:\\Program Files\\Android\\sdk\\platform-tools\\adb.exe',
        'C:\\Program Files (x86)\\Android\\sdk\\platform-tools\\adb.exe',
        # 用户自定义安装路径
        os.path.expanduser('~/android-sdk/platform-tools/adb'),
        os.path.expanduser('~/Android/Sdk/platform-tools/adb')
    ]

    for path in common_paths:
        if os.path.exists(path) and os.access(path, os.X_OK):
            return path

    return None


class QNXProcessMonitor:
    def __init__(self, process_name, interval=5, output_file=None, device_type="device1",
                 ip1="192.168.8.1", username1="root", password1="sv2655888",
                 ip2="172.31.207.204", username2="root", adb_path=None):
        """
        初始化QNX进程监控器

        Args:
            process_name: 要监控的进程名
            interval: 监控间隔时间(秒)
            output_file: 输出CSV文件名(可选)
            device_type: 设备类型("device1"或"device2")
            ip1: 设备1的QNX系统IP地址
            username1: 设备1的登录用户名
            password1: 设备1的登录密码
            ip2: 设备2的QNX系统IP地址
            username2: 设备2的登录用户名
            adb_path: adb命令路径(可选)
        """
        self.process_name = process_name
        self.interval = interval
        self.device_type = device_type
        self.ip1 = ip1
        self.username1 = username1
        self.password1 = password1
        self.ip2 = ip2
        self.username2 = username2
        self.adb_path = adb_path

        # 设置输出文件
        if output_file:
            self.output_file = output_file
        else:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            self.output_file = f"{process_name}_monitor_{timestamp}.csv"

        # 创建CSV文件并写入表头
        with open(self.output_file, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['时间戳', '进程名', 'PID', 'CPU使用率(%)', '总内存(KB)', '堆内存(KB)'])

        print(f"监控数据将保存到: {self.output_file}")

    def connect_to_qnx_device1(self):
        """连接到需要密码验证的QNX系统并返回终端会话"""
        # 确保adb路径有效
        if not self.adb_path:
            raise ValueError("无法找到adb命令。请确保已安装adb并添加到PATH环境变量中，或在脚本中设置ADB_PATH。")

        # 启动ADB shell
        print(f"使用adb路径: {self.adb_path}")
        try:
            if DEBUG_MODE:
                print("DEBUG: 尝试启动adb shell...")
            adb_shell = pexpect.spawn(f'{self.adb_path} shell', timeout=15)
            if DEBUG_MODE:
                print("DEBUG: adb shell已启动")
        except Exception as e:
            if IS_WINDOWS:
                error_msg = f"启动adb shell失败: {str(e)}。请检查adb是否正确安装，设备是否连接。"
            else:
                error_msg = f"启动adb shell失败: {str(e)}。请检查adb是否正确安装，设备是否连接。"
            if DEBUG_MODE:
                print(f"DEBUG: {error_msg}")
            raise RuntimeError(error_msg)

        if DEBUG_MODE:
            print("DEBUG: 等待adb shell提示符...")
        index = adb_shell.expect(['#', '>', pexpect.TIMEOUT, pexpect.EOF], timeout=15)
        if index == 2:
            if DEBUG_MODE:
                print("DEBUG: adb shell连接超时")
            raise TimeoutError("adb shell连接超时。请检查设备是否已连接并开启USB调试。")
        elif index == 3:
            if DEBUG_MODE:
                print("DEBUG: adb shell连接失败")
            raise ConnectionError("adb shell连接失败。请检查设备是否已连接并开启USB调试。")

        if DEBUG_MODE:
            print(f"DEBUG: adb shell已连接，提示符索引: {index}")

        # 启动telnet连接
        print(f"连接到QNX系统 {self.ip1}...")
        if DEBUG_MODE:
            print(f"DEBUG: 发送命令: busybox telnet {self.ip1}")
        adb_shell.sendline(f'busybox telnet {self.ip1}')

        # 等待登录提示
        if DEBUG_MODE:
            print("DEBUG: 等待telnet登录提示...")
        index = adb_shell.expect(['login:', 'Connection refused', pexpect.TIMEOUT], timeout=15)
        if index == 1:
            if DEBUG_MODE:
                print("DEBUG: telnet连接被拒绝")
            raise ConnectionError(f"telnet连接被拒绝。请检查IP地址 {self.ip1} 和网络连接。")
        elif index == 2:
            if DEBUG_MODE:
                print("DEBUG: telnet连接超时")
            raise TimeoutError(f"telnet连接超时。请检查设备 {self.ip1} 是否可达。")

        if DEBUG_MODE:
            print("DEBUG: 收到登录提示")

        # 登录
        print(f"使用用户名 {self.username1} 登录...")
        adb_shell.sendline(self.username1)
        if DEBUG_MODE:
            print("DEBUG: 等待密码提示...")
        adb_shell.expect('Password:', timeout=15)
        if DEBUG_MODE:
            print("DEBUG: 收到密码提示，发送密码...")
        adb_shell.sendline(self.password1)

        # 确认登录成功
        if DEBUG_MODE:
            print("DEBUG: 等待登录确认...")
        index = adb_shell.expect(['#', '>', 'Login incorrect', pexpect.TIMEOUT], timeout=10)
        if index == 2:
            if DEBUG_MODE:
                print("DEBUG: 登录失败")
            raise ValueError("登录失败。请检查用户名和密码。")
        elif index == 3:
            if DEBUG_MODE:
                print("DEBUG: 登录超时")
            raise TimeoutError("登录后无响应。请检查设备状态。")

        if DEBUG_MODE:
            print(f"DEBUG: 登录成功，提示符索引: {index}")

        print("成功连接到QNX系统")
        return adb_shell

    def connect_to_qnx_device2(self):
        """连接到只需账号不需密码的QNX系统并返回终端会话"""
        # 确保adb路径有效
        if not self.adb_path:
            raise ValueError("无法找到adb命令。请确保已安装adb并添加到PATH环境变量中，或在脚本中设置ADB_PATH。")

        # 首先执行adb root
        print(f"执行adb root...")
        try:
            if DEBUG_MODE:
                print("DEBUG: 尝试执行adb root...")
            root_process = pexpect.spawn(f'{self.adb_path} root', timeout=10)
            root_process.expect([pexpect.EOF, pexpect.TIMEOUT], timeout=10)
            root_process.close()
            # 给设备一些时间重启adbd
            if DEBUG_MODE:
                print("DEBUG: adb root执行完成，等待设备重启adbd...")
            time.sleep(3)
        except Exception as e:
            if DEBUG_MODE:
                print(f"DEBUG: adb root执行失败: {str(e)}")
            raise RuntimeError(f"执行adb root失败: {str(e)}。请检查设备是否支持root访问。")

        # 启动ADB shell
        print(f"启动adb shell...")
        try:
            if DEBUG_MODE:
                print("DEBUG: 尝试启动adb shell...")
            adb_shell = pexpect.spawn(f'{self.adb_path} shell', timeout=15)
            if DEBUG_MODE:
                print("DEBUG: adb shell已启动")
        except Exception as e:
            if DEBUG_MODE:
                print(f"DEBUG: adb shell启动失败: {str(e)}")
            raise RuntimeError(f"启动adb shell失败: {str(e)}。请检查adb是否正确安装，设备是否连接。")

        if DEBUG_MODE:
            print("DEBUG: 等待adb shell提示符...")
        index = adb_shell.expect(['#', '>', pexpect.TIMEOUT, pexpect.EOF], timeout=15)
        if index == 2:
            if DEBUG_MODE:
                print("DEBUG: adb shell连接超时")
            raise TimeoutError("adb shell连接超时。请检查设备是否已连接并开启USB调试。")
        elif index == 3:
            if DEBUG_MODE:
                print("DEBUG: adb shell连接失败")
            raise ConnectionError("adb shell连接失败。请检查设备是否已连接并开启USB调试。")

        if DEBUG_MODE:
            print(f"DEBUG: adb shell已连接，提示符索引: {index}")

        # 启动telnet连接
        print(f"连接到QNX系统 {self.ip2}...")
        if DEBUG_MODE:
            print(f"DEBUG: 发送命令: busybox telnet {self.ip2}")
        adb_shell.sendline(f'busybox telnet {self.ip2}')

        # 等待登录提示
        if DEBUG_MODE:
            print("DEBUG: 等待telnet登录提示...")
        index = adb_shell.expect(['login:', 'Connection refused', pexpect.TIMEOUT], timeout=15)
        if index == 1:
            if DEBUG_MODE:
                print("DEBUG: telnet连接被拒绝")
            raise ConnectionError(f"telnet连接被拒绝。请检查IP地址 {self.ip2} 和网络连接。")
        elif index == 2:
            if DEBUG_MODE:
                print("DEBUG: telnet连接超时")
            raise TimeoutError(f"telnet连接超时。请检查设备 {self.ip2} 是否可达。")

        if DEBUG_MODE:
            print("DEBUG: 收到登录提示")

        # 登录(只需要用户名，不需要密码)
        print(f"使用用户名 {self.username2} 登录...")
        adb_shell.sendline(self.username2)

        # 确认登录成功，这里不会提示输入密码
        if DEBUG_MODE:
            print("DEBUG: 等待登录确认...")
        index = adb_shell.expect(['#', '>', 'Login incorrect', pexpect.TIMEOUT], timeout=10)
        if index == 2:
            if DEBUG_MODE:
                print("DEBUG: 登录失败")
            raise ValueError("登录失败。请检查用户名。")
        elif index == 3:
            if DEBUG_MODE:
                print("DEBUG: 登录超时")
            raise TimeoutError("登录后无响应。请检查设备状态。")

        if DEBUG_MODE:
            print(f"DEBUG: 登录成功，提示符索引: {index}")

        print("成功连接到QNX系统")
        return adb_shell

    def connect_to_qnx(self):
        """根据设备类型选择适当的连接方法"""
        if self.device_type == "device1":
            return self.connect_to_qnx_device1()
        elif self.device_type == "device2":
            return self.connect_to_qnx_device2()
        else:
            raise ValueError(f"不支持的设备类型: {self.device_type}。请使用 'device1' 或 'device2'。")

    def get_process_info(self, terminal):
        """获取指定进程的CPU和内存信息"""
        # 获取进程的PID和CPU使用率
        terminal.sendline('top -i 1 -d')
        terminal.expect(['#', '>'])

        # 根据操作系统处理不同的输出格式
        if IS_WINDOWS:
            # Windows上wexpect返回的是字符串，不需要decode
            top_output = terminal.before
        else:
            # 非Windows系统上pexpect返回的是bytes，需要decode
            top_output = terminal.before.decode('utf-8', errors='ignore')

        # 根据示例数据调整正则表达式
        # 示例格式: 181600357     2  10 Sem       1:01:32   1.69% qvm
        pid_pattern = re.compile(r'(\d+)\s+\d+\s+\d+\s+\w+\s+[\d:]+\s+([\d.]+)%\s+' + re.escape(self.process_name))

        # 在整个top输出中查找进程信息
        pid_match = None
        for line in top_output.splitlines():
            match = pid_pattern.search(line)
            if match and self.process_name in line:
                pid_match = match
                break

        if not pid_match:
            print(f"警告: 未找到进程 '{self.process_name}'")
            return None

        pid = pid_match.group(1)
        cpu_usage = pid_match.group(2)

        # 获取内存信息
        terminal.sendline(f'showmem -p {pid}')
        terminal.expect(['#', '>'])

        # 根据操作系统处理不同的输出格式
        if IS_WINDOWS:
            # Windows上wexpect返回的是字符串，不需要decode
            mem_output = terminal.before
        else:
            # 非Windows系统上pexpect返回的是bytes，需要decode
            mem_output = terminal.before.decode('utf-8', errors='ignore')

        # 根据示例数据调整正则表达式
        # 示例格式: logservice |        2326601 |          29444 |          30328 |            128 |           3064 |            428 |          17936 |                 0 |           8772 |
        mem_pattern = re.compile(r'\s+' + re.escape(self.process_name) + r'\s+\|\s+' + re.escape(
            pid) + r'\s+\|\s+\d+\s+\|\s+(\d+)\s+\|\s+\d+\s+\|\s+\d+\s+\|\s+\d+\s+\|\s+(\d+)\s+\|')

        mem_match = None
        for line in mem_output.splitlines():
            match = mem_pattern.search(line)
            if match:
                mem_match = match
                break

        if not mem_match:
            # 尝试使用一个更宽松的正则表达式
            alt_mem_pattern = re.compile(
                r'\s+\S+\s+\|\s+' + re.escape(pid) + r'\s+\|[^|]+\|\s+(\d+)\s+\|[^|]+\|[^|]+\|[^|]+\|\s+(\d+)\s+\|')

            for line in mem_output.splitlines():
                match = alt_mem_pattern.search(line)
                if match:
                    mem_match = match
                    break

        if not mem_match:
            # 如果仍然未匹配，则尝试匹配最后一行的总内存使用
            total_mem_pattern = re.compile(r'Total process memory in use : (\d+)\(KB\)')
            total_match = total_mem_pattern.search(mem_output)

            if total_match:
                total_mem = total_match.group(1)
                heap_mem = "N/A"
            else:
                print(f"警告: 无法获取进程 '{self.process_name}' (PID: {pid}) 的内存信息")
                total_mem = "N/A"
                heap_mem = "N/A"
        else:
            total_mem = mem_match.group(1)
            heap_mem = mem_match.group(2)

        return {
            'pid': pid,
            'cpu_usage': cpu_usage,
            'total_mem': total_mem,
            'heap_mem': heap_mem
        }

    def start_monitoring(self):
        """开始持续监控进程"""
        try:
            # 连接到QNX系统
            terminal = self.connect_to_qnx()

            print(f"开始监控进程 '{self.process_name}'，每 {self.interval} 秒采集一次数据...")
            print("按 Ctrl+C 停止监控")

            while True:
                # 获取进程信息
                process_info = self.get_process_info(terminal)

                if process_info:
                    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

                    # 打印到控制台
                    print(f"[{timestamp}] 进程: {self.process_name}, PID: {process_info['pid']}, "
                          f"CPU: {process_info['cpu_usage']}%, 总内存: {process_info['total_mem']}KB, "
                          f"堆内存: {process_info['heap_mem']}KB")

                    # 写入CSV文件
                    with open(self.output_file, 'a', newline='') as csvfile:
                        writer = csv.writer(csvfile)
                        writer.writerow([
                            timestamp,
                            self.process_name,
                            process_info['pid'],
                            process_info['cpu_usage'],
                            process_info['total_mem'],
                            process_info['heap_mem']
                        ])

                # 等待指定的间隔时间
                time.sleep(self.interval)

        except KeyboardInterrupt:
            print("\n监控已停止")
        except Exception as e:
            print(f"发生错误: {str(e)}")
        finally:
            if 'terminal' in locals():
                terminal.sendline('exit')
                terminal.close()
            print(f"监控数据已保存到: {self.output_file}")


def main():
    """主函数，使用内部配置参数启动监控"""
    try:
        # 查找adb路径
        adb_path = find_adb_path()
        if not adb_path:
            print("错误: 无法找到adb命令。")
            print("请尝试以下解决方案:")
            print("1. 确保已安装Android Debug Bridge (adb)")
            print("2. 将adb添加到系统PATH环境变量中")
            print("3. 在脚本配置区域设置ADB_PATH变量为adb的完整路径")
            print("\n可能的adb路径位置:")
            print("- macOS: ~/Library/Android/sdk/platform-tools/adb")
            print("- Linux: /usr/bin/adb 或 /opt/android-sdk/platform-tools/adb")
            print("- Windows: C:\\Program Files\\Android\\sdk\\platform-tools\\adb.exe")
            return

        # 根据设备类型设置参数
        if DEVICE_TYPE == "device1":
            print("使用设备类型1: 需要密码验证的设备")
            ip = DEVICE1_QNX_IP
            username = DEVICE1_USERNAME
            password = DEVICE1_PASSWORD
        elif DEVICE_TYPE == "device2":
            print("使用设备类型2: 只需账号不需密码的设备")
            ip = DEVICE2_QNX_IP
            username = DEVICE2_USERNAME
            password = None
        else:
            print(f"错误: 不支持的设备类型 '{DEVICE_TYPE}'。请使用 'device1' 或 'device2'。")
            return

        monitor = QNXProcessMonitor(
            process_name=PROCESS_NAME,
            interval=INTERVAL,
            output_file=OUTPUT_FILE,
            device_type=DEVICE_TYPE,
            ip1=DEVICE1_QNX_IP,
            username1=DEVICE1_USERNAME,
            password1=DEVICE1_PASSWORD,
            ip2=DEVICE2_QNX_IP,
            username2=DEVICE2_USERNAME,
            adb_path=adb_path
        )

        monitor.start_monitoring()
    except Exception as e:
        print(f"启动监控失败: {str(e)}")


if __name__ == '__main__':
    main()